// package pktbuilder 用于构造自定义的网络数据包
package pktbuilder

import (
	"bytes"
	"log"
	"net"
	"os"

	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
)

var (
	// CustomEthType
	CustomEthType layers.EthernetType = 0xF8F8

	// CustomIPProtocol
	CustomIPProtocol layers.IPProtocol = 0xF8

	// CustomIPPayload
	CustomPayload = bytes.Repeat([]byte{0x88}, 50)
)

// serializableLayerList
type serializableLayerList struct {
	Type              uint16 // type or proto
	SerializableLayer gopacket.SerializableLayer
	next              *serializableLayerList
}

// BuildingEtherPkt 构建以太网包字节
func BuildingEtherPkt(srcMAC, dstMAC net.HardwareAddr, next *serializableLayerList) []byte {
	var (
		buffer           = gopacket.NewSerializeBuffer()
		serializeOptions = gopacket.SerializeOptions{ComputeChecksums: true, FixLengths: true}
	)

	var root = &serializableLayerList{}
	eth := &layers.Ethernet{SrcMAC: srcMAC, DstMAC: dstMAC}
	if next != nil {
		eth.EthernetType = layers.EthernetType(next.Type)
		root = &serializableLayerList{SerializableLayer: eth, next: next}

	} else {
		eth.EthernetType = CustomEthType
		root = &serializableLayerList{SerializableLayer: eth, next: &serializableLayerList{SerializableLayer: gopacket.Payload(CustomPayload)}}
	}

	serializableLayers := make([]gopacket.SerializableLayer, 0, 10)
	for {
		serializableLayers = append(serializableLayers, root.SerializableLayer)
		root = root.next
		if root == nil {
			break
		}
	}

	for i := len(serializableLayers) - 1; i >= 0; i-- {
		layer := serializableLayers[i]
		err := serializableLayers[i].SerializeTo(buffer, serializeOptions)
		if err != nil {
			log.Println(err)
			return nil
		}
		buffer.PushLayer(layer.LayerType())
	}

	return buffer.Bytes()
}

// GenLayerThatDot1Q 创建 VLAN 层
func GenLayerThatDot1Q(pri uint8, vid uint16, nextLayer *serializableLayerList) *serializableLayerList {
	var (
		thisLayer = &serializableLayerList{Type: uint16(layers.EthernetTypeDot1Q)}
		dot1q     = &layers.Dot1Q{Priority: pri, VLANIdentifier: vid}
	)

	if nextLayer != nil {
		dot1q.Type = layers.EthernetType(nextLayer.Type)
		thisLayer.next = nextLayer
	} else {
		dot1q.Type = CustomEthType
		thisLayer.next = &serializableLayerList{SerializableLayer: gopacket.Payload(CustomPayload)}
	}

	thisLayer.SerializableLayer = dot1q
	return thisLayer
}

// GenLayerThatIP4 创建 IPv4 层
func GenLayerThatIP4(srcIP, dstIP net.IP, nextLayer *serializableLayerList) *serializableLayerList {
	var (
		thisLayer = &serializableLayerList{Type: uint16(layers.EthernetTypeIPv4)}
		ip4       = &layers.IPv4{
			Version: 4, Id: uint16(os.Getpid() & 0xffff), TTL: 32,
			SrcIP: srcIP, DstIP: dstIP,
		}
	)

	type canSetNetLayer interface {
		SetNetworkLayerForChecksum(gopacket.NetworkLayer) error
	}

	if nextLayer != nil {
		ip4.Protocol = layers.IPProtocol(nextLayer.Type)
		thisLayer.next = nextLayer

		if set, ok := nextLayer.SerializableLayer.(canSetNetLayer); ok {
			_ = set.SetNetworkLayerForChecksum(ip4)
		}
	} else {
		ip4.Protocol = CustomIPProtocol
		thisLayer.next = &serializableLayerList{SerializableLayer: gopacket.Payload(CustomPayload)}
	}

	thisLayer.SerializableLayer = ip4
	return thisLayer
}

// GenLayerThatUDP 创建 TCP 层
func GenLayerThatUDP(srcPort, dstPort uint16, payload []byte) *serializableLayerList {
	var (
		thisLayer = &serializableLayerList{Type: uint16(layers.IPProtocolUDP)}
		udp       = &layers.UDP{SrcPort: layers.UDPPort(srcPort), DstPort: layers.UDPPort(dstPort)}
	)

	if payload != nil {
		thisLayer.next = &serializableLayerList{SerializableLayer: gopacket.Payload(payload)}
	} else {
		thisLayer.next = &serializableLayerList{SerializableLayer: gopacket.Payload(CustomPayload)}
	}

	thisLayer.SerializableLayer = udp
	return thisLayer
}

// GenLayerThatTCP 创建 TCP 层
func GenLayerThatTCP(srcPort, dstPort uint16, seq uint32, payload []byte) *serializableLayerList {
	var (
		thisLayer = &serializableLayerList{Type: uint16(layers.IPProtocolTCP)}
		tcp       = &layers.TCP{
			SrcPort: layers.TCPPort(srcPort), DstPort: layers.TCPPort(dstPort), Seq: seq,
			Window: 512, PSH: true, ACK: true,
		}
	)

	if payload != nil {
		thisLayer.next = &serializableLayerList{SerializableLayer: gopacket.Payload(payload)}
	} else {
		thisLayer.next = &serializableLayerList{SerializableLayer: gopacket.Payload(CustomPayload)}
	}

	thisLayer.SerializableLayer = tcp
	return thisLayer
}

// GenOtherLayer
func GenOtherLayer(typ uint16, serializableLayer gopacket.SerializableLayer) *serializableLayerList {
	return &serializableLayerList{Type: typ, SerializableLayer: serializableLayer}
}
